(* module which implements a context graph data structure, including it's supported operations *)
module ContextGraph
    open System.Collections.Generic
    open BIT_ARRAY
    (* this "global" variable stores the number of hash table buckets in graphs created through the newGraph() function.
        in order to improve performance - it should be initialized to at least 2^X, where X is the number of variables in the analysed program. *)
    let numBuckets = ref 100;

    (* a flag indicating if edges with empty DFRs should be added to graph. *)
    let keepEmptyEdges = ref false;

    let debug = ref true;

    let debugPrint(s) = if (!debug) then (printfn "%s" s) else ();

    type Mode = Normal | Label of string

    type Context = { resetCtx: Bitarray; modeCtx: Mode; constantCtx: Bitarray }

    (* the data which is hanging on the context graph edges *)
    type edgeTag = setOfPairs.DFR ref

    (*the graph type - it is actually a hash table with BitArray keys *)
    type graph = Dictionary<Context, Dictionary<Context, edgeTag>>

    (* a tagged or untagged edge datatype *)
    type edge = E of Context * Context * edgeTag | UE of Context * Context

    (* returns the specified edge's source context *)
    let edgeSrc = function
        | E(p, _, _) -> p
        | UE(p, _)  -> p

    (* returns the specified edge's target context *)
    let edgeTgt = function
        | E(_, q, _) -> q
        | UE(_, q)  -> q

    (* If dfr is non-empty:
       1. edge exists in g - updates exiting edge's DFR to be: LUB(existing_DFR,  new_DFR) .
       2. Otherwise, adds the new edge to g. *)
    let addEdge(g : graph, e) =
        match e with
        | E(p, q, dfr) ->
            if ( !keepEmptyEdges || (not (setOfPairs.dfrEmpty(!dfr)))) then
                if (g.ContainsKey(p)) then
                    if g.[p].ContainsKey(q) then
                        let dfr' = g.[p].[q]
                        g.[p].[q] := setOfPairs.sqcup(!dfr',!dfr)
                        (* in case edge exist in graph update edge to be LUB(existing_DFR, new_DFR) *)
                    else
                        g.[p].Add(q, dfr)
                        (* Otherwise - insert new edge to g. *)
                else
                    g.Add(p,new Dictionary<Context, edgeTag>())
                    g.[p].Add(q, dfr)
            else ()
        | _ -> failwith "ContextGraph.addEdge: invalid argument"

    (* inserts a list of edges to the graph g.
        If a specified edge is already in g, the new edge's DFR will be joined with the existing edge's DFR *)
    let rec addEdges = function
        | (g : graph, []) -> ()
        | (g : graph, e::el) -> (addEdge(g,e); addEdges(g,el));

    let removeEdge(g : graph, e : edge) : bool =
        match e with
        | UE(p, q) ->
            if g.ContainsKey(p) && g.[p].ContainsKey(q) then
                let _ = g.[p].Remove(q)
                let _ = if g.[p].Count = 0 then g.Remove(p) else false
                true
            else false
        | _ -> failwith "ContextGraph.removeEdge: invalid argument"

    let removeEdgesByPreContext(g : graph ref, pc : Context) = (!g).Remove(pc)

      (* constructor for new graphs - creates a new graph given the specified edge list *)
    let newGraph = function
        | [] -> new graph()
        | edgeList ->
            let g = new graph()
            addEdges(g, edgeList)
            g

    (* returns the edge tag of the specified edge *)
    let getEdgeTag(g : graph, e) =
        match e with
        | UE(p, q) -> if g.ContainsKey(p) && g.[p].ContainsKey(q) then
                                            Some g.[p].[q]
                                          else
                                            None
        | _ -> failwith "ContextGraph.getEdgeTag: invalid argument"

    (* returns a list of all edges in graph - not including edge tags (UE edge constructor is used to create each edge) *)
    let untaggedEdges(g : graph) =
        let pUntaggedEdges(p : Context, tgt : Dictionary<Context,edgeTag>) =
            let tgtContexts = (Seq.toList tgt.Keys)
            List.map (fun (q : Context) -> UE(p, q)) tgtContexts
        let ps = List.map (fun c -> (c,g.[c])) (Seq.toList(g.Keys))
        List.concat (List.map pUntaggedEdges ps)

    (* returns all edges in graph including edge tags (E edge constructor is usedto create each edge) *)
    let getEdges(g : graph) =
        let pTaggedEdges(p : Context,tgt : Dictionary<Context,edgeTag>) =
              let tgtContextAndDfrPairs = (Seq.toList tgt)
              let pKeyValue2edge(pair : KeyValuePair<Context,edgeTag> ) =
                    E(p, pair.Key, pair.Value)
              List.map pKeyValue2edge tgtContextAndDfrPairs

        let ps = List.map (fun c -> (c,g.[c])) (Seq.toList(g.Keys))
        List.concat (List.map pTaggedEdges ps)

    (* creates an empty context which can hold up to the specified number of variables *)
    let makeEmptyContext(size : int) : Context = { resetCtx = Bitarray.create(size);  modeCtx = Mode.Normal; constantCtx = Bitarray.create(size) }

    (* creates a context which can hold up to the specified number of variables and is filled with the specified list of variable indices *)
    let makeContext(size : int, vars : int list) : Context =
        let c = makeEmptyContext(size)
        List.iter (fun i->(Bitarray.set (i-1) c.resetCtx)) vars
        c

    (* returns a copy of the srcContext, assuming that numVars is the number of variables in srcContext *)
    let copyContext(srcContext) = { resetCtx = Bitarray.copy srcContext.resetCtx; modeCtx = srcContext.modeCtx; constantCtx = Bitarray.copy srcContext.constantCtx }

    (* adds variable with index x to reset context c *)
    let addVarToResetContext x c = Bitarray.set (x-1) c.resetCtx

    (* removes variable with index x to reset context c *)
    let removeVarFromResetContext x  c = Bitarray.unset (x-1) c.resetCtx

    (* if 'add' is true - adds x to reset context c, otherwise removes x from reset context c *)
    let setVarInResetContext x c add = if add then addVarToResetContext x c else removeVarFromResetContext x c

    (* adds variable with index x to constant context c *)
    let addVarToConstantContext x c = Bitarray.set (x-1) c.constantCtx

    (* removes variable with index x to constant context c *)
    let removeVarFromConstantContext x  c = Bitarray.unset (x-1) c.constantCtx

    (* if 'add' is true - adds x to constant context c, otherwise removes x from constant context c *)
    let setVarInConstantContext x c add = if add then addVarToConstantContext x c else removeVarFromConstantContext x c

    (* returns a list of variable indices which are contained in the specified context. *)
    let resetContextToVarList(c : Context) : int list =
        List. map (fun i -> i+1) (Bitarray.setBits c.resetCtx)
    let ConstantContextToVarList(c : Context) : int list =
        List. map (fun i -> i+1) (Bitarray.setBits c.constantCtx)

    let prtMode(mode) = 
      match mode with
        | Label lbl -> lbl
        | _ -> "Normal"

     (* outputs the context's variables to a nicely formatted string. *)
    let prtContext(c : Context) =
        let prtVar(index, comma, array) = 
            if (Bitarray.isSet index array) then
                (if (comma) then ", " else "") + Parser.parserLastVarToIndexMap().Item(index + 1)
            else ""

        let rec prtVars(index, len, first, array) =
            if (index < len) then
                  let vr = prtVar(index, not first, array)
                  let first1 = first && ((String.length vr) = 0)
                  vr + prtVars(index+1, len, first1, array)
            else ""

        ("{R: " + prtVars(0, (Bitarray.length c.resetCtx), true, c.resetCtx) + "} " 
        + " {C: " + prtVars(0, (Bitarray.length c.constantCtx), true, c.constantCtx) + "} "
        + prtMode(c.modeCtx))

    let rec prtContexts = function
        | []   -> ""
        | c::t -> prtContext(c) + ", " + prtContexts(t)

    let inResetContext(v : int, c : Context) = Bitarray.isSet (v-1) c.resetCtx
    let inConstantContext(v : int, c : Context) = Bitarray.isSet (v-1) c.constantCtx

    let contextToNum(c : Context) =
        let rec computeNum = function
            | (numSoFar : int, powerOf2 : int, 0) -> numSoFar
            | (numSoFar : int, powerOf2, index : int) -> computeNum(
                                                           (if (Bitarray.isSet index c.resetCtx) then numSoFar+powerOf2 else 0),
                                                           powerOf2 * 2,
                                                           index-1)

        computeNum(0,1,(Bitarray.length c.resetCtx)-1)

    (* a partial order on edges, based on lexicographic ordering of the source and target contexts hash code. *)
    let edgeSmaller(e1,e2) =
        let p = contextToNum(edgeSrc(e1))
        let p' = contextToNum(edgeSrc(e2))
        p < p' || (p = p' && (contextToNum(edgeTgt(e1)) < contextToNum(edgeTgt(e2))))

    (* inserts an edge into a sorted edge list *)
    let rec insertEdge = function
        | (e,[]) -> [e]
        | (e,((h::t) as el)) -> if (edgeSmaller(h,e)) then
                                  h::insertEdge(e,t)
                                else
                                  e::el
    (* sort an edge list by lexical order <src_context, tgt_context> *)
    let rec sortEdges = function
        | [] -> []
        | e::t -> insertEdge(e,sortEdges(t));

    (* returns true iff context c1 is equal to context c2. *)
    let eqContexts(c1 : Context, c2 : Context) = 
        Bitarray.areAllUnset(c1.resetCtx ^^^ c2.resetCtx) && 
        c1.modeCtx = c2.modeCtx && 
        Bitarray.areAllUnset(c1.constantCtx ^^^ c2.constantCtx)

    let contextSmaller(c1,c2) =
        let p = contextToNum(c1)
        let p' = contextToNum(c2)
        p < p'

    (* inserts a context into a sorted context list *)
    let rec insertContext c cl =
        match cl with
        | []           -> [c]
        | (h::t) as cl ->
            if (contextSmaller(h,c)) then
                h::(insertContext c t)
            else c::cl

    let sortContexts(contextList : Context list) =
        List.sort contextList

        (* returns all pre contexts of graph g. *)
    let preContexts(g : graph) : Context list =
        List.map (fun e -> match e with
                           |UE(p, _)-> p
                           | _ -> failwith "ContextGraph.preContexts: Invalid argument.")
                 (untaggedEdges(g))

    (* returns true if every context of preContextsList is a pre-contexts of graph g and vice-versa. *)
    let preContextsEqual(preContextsList, g) =
        let rec equalContextLists = function
            | ([],[]) -> true
            | (_,[]) -> false
            | ([],_) -> false
            | (c1::cl1, c2::cl2) -> eqContexts(c1,c2) && equalContextLists(cl1,cl2)
        equalContextLists(preContexts(g),preContextsList)

    (* returns true if every context of preContextsList is a pre-contexts of graph g. *)
    let preContextsIn(preContextsLis, g : graph) =
        let yes = ref true;
        for c in preContextsLis do
            yes := !yes && g.ContainsKey(c)
        !yes

    let prtDfr((setOfPairs.Dfr(M,R)) as d) = "\n    M=[ " + M.ToString() + "]\n    R=[ " + setOfPairs.prtR(d) + "]";

    (* outputs the edge context variables and DFR to a nicely formatted string. *)
    let prtEdge = function
        | E(p, q, dfr) ->
          match dfr with
            | dfrRef ->
              match !dfrRef with
                | setOfPairs.Dfr(M, R) as d -> prtContext(p) + "-->" + prtContext(q) + " " + prtDfr(d)
        | UE(p, q) -> prtContext(p) + " ? "  + prtContext(q)

    (* for all edge in the list - outputs the edge context variables and DFR to a nicely formatted string. *)
    let rec prtEdgeList = function
        | ([]) -> ""
        | (e::t) -> prtEdge(e) + "\n" + prtEdgeList(t);

     (* for all edge in the graph - outputs the edge context variables and DFR to a nicely formatted string. *)
    let prtEdges(g : graph) =
        prtEdgeList(sortEdges(getEdges(g)))

    (* for all edge in the graph - outputs the edge context variables and DFR to a nicely formatted string. *)
    let printEdges(g : graph) =
        printfn "%s" (prtEdgeList(sortEdges(getEdges(g))))

    (* returns all post contexts of graph g. *)
    let postContexts(g : graph) : Context list = List.map edgeTgt (untaggedEdges(g));

    (* returns true if g1 and g2 are equal in all aspects *)
    let equal(g1 : graph, g2 : graph) : bool =
        let res = ref true
        let mutable e = g1.Keys.GetEnumerator()
        let () = while e.MoveNext() && !res do
                    let p = e.Current
                    if (g2.ContainsKey(p)) then
                        let qds = g1.[p]
                        let qds' = g2.[p]
                        let mutable e1 = qds.Keys.GetEnumerator()
                        while e1.MoveNext() && !res do
                            let q = e1.Current
                            if qds'.ContainsKey(q) then
                                res := !res && (setOfPairs.dfrsame(!qds.[q],!qds'.[q]))
                            else res := false
                    else
                        res := false
        !res

    (* returns true if c is a context which contains all variables (all bits are 1). *)
    let isFullContext(c : Context) = Bitarray.areAllSet c.resetCtx

    (* computes all possible contexts (except the full context) over a set of variables,
       where highestVar is the highest variable index *)
    let allContexts(highestVar) =
        let rec intToBinaryInternal = function
            | (0,n,k) -> []
            | (i,n,k) -> if ((i % 2) = 1) then
                            (n-k+1)::intToBinaryInternal(i / 2,n,k-1)
                         else
                            intToBinaryInternal(i / 2,n,k-1) ;
        let intToBinaryList(n) = intToBinaryInternal(n,n,n);

        let rec makeAllContexts(highest_var,i) =
            let c = makeContext(highest_var, intToBinaryList(i));
            if (isFullContext(c)) then [] else  c :: makeAllContexts(highest_var,i+1)

        let contexts = makeAllContexts(highestVar, 0);
        let _ = debugPrint("All contexts: " + prtContexts(contexts) + "\n");
        contexts

    let compose (g1 : graph) (g2 : graph) prodFucntion =
        let g = newGraph([])
        let () = for g1PreCtx in g1.Keys do
                    for g1Edge in g1.[g1PreCtx] do
                        let g1PostCtx = g1Edge.Key
                        let found, g2Edges = g2.TryGetValue(g1PostCtx)
                        if (found) then
                            let g1Dfr = g1Edge.Value
                            for g2Edge in g2Edges do
                                let g2PostCtx = g2Edge.Key
                                let g2Dfr = g2Edge.Value
                                let dep = setOfPairs.prod !g1Dfr !g2Dfr prodFucntion
                                addEdge(g, E(g1PreCtx, g2PostCtx, (ref dep)))
                        else ()
        g

    let debugPrintEdges(edges') =
        if (!debug) then debugPrint(prtEdgeList(edges')) else ()

    (* computes the join (LUB) of two graphs *)
    let joinAsLeastUpperBound (g1 : graph) (g2 : graph) =
        newGraph( ((getEdges(g1)) @ (getEdges(g2))) )

    let clone(g : graph) : graph =
        newGraph(getEdges(g));

    (* removes from g all edges which do not exit a context from pc. *)
    let filterEdgesByPreContexts(g : graph ref, pc : Context list) =
        let pcs = Set.ofList pc
        let pcg = preContexts(!g)
        for p in pcg do
            if not (pcs.Contains p) then
                let _ = removeEdgesByPreContext(g, p)
                ()
            else ()

 (* module ContextGraph *)